from langchain.embeddings.dashscope import DashScopeEmbeddings
from langchain_chroma import Chroma
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader, WebBaseLoader
from dotenv import find_dotenv, load_dotenv
from langchain.chains import RetrievalQA
import bs4
import os
from langchain_community.chat_models.tongyi import ChatTongyi

load_dotenv(find_dotenv())
DASHSCOPE_API_KEY = os.environ["DASHSCOPE_API_KEY"]


def init_index():
  loader = DirectoryLoader('../text', glob='**/jianli1.txt')
  documents = loader.load()
  text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
  split_docs = text_splitter.split_documents(documents)
  # 初始化 embeddings 对象
  embeddings = DashScopeEmbeddings(
      model="text-embedding-v1",
  )
  vector = Chroma(collection_name='customer', embedding_function=embeddings,
                  persist_directory='../chroma')
  vector.add_documents(documents=split_docs, embeddings=embeddings)


def web_search():
  # 从浏览器获取数据
  loader = WebBaseLoader(
      web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
      bs_kwargs=dict(
          parse_only=bs4.SoupStrainer(
              class_=("post-content", "post-title", "post-header")
          )
      ),
  )
  # loader = WebBaseLoader(
  #     web_path="https://www.360wenmi.com/f/filecx3o1e5r.html",
  #     # header_template = None,
  #     # verify_ssl = True,
  #     # proxies = None,
  #     # continue_on_failure = False,
  #     # autoset_encoding = True,
  #     # encoding = None,
  #     # web_paths = (),
  #     # requests_per_second = 2,
  #     # default_parser = "html.parser",
  #     # requests_kwargs = None,
  #     # raise_for_status = False,
  #     # bs_get_text_kwargs = None,
  #     # bs_kwargs = None,
  #     # session = None,
  #     show_progress=True
  # )
  documents = loader.load()
  text_splitter = CharacterTextSplitter(chunk_size=10000, chunk_overlap=0)
  split_docs = text_splitter.split_documents(documents)
  # 初始化 embeddings 对象
  embeddings = DashScopeEmbeddings(
      model="text-embedding-v1",
  )
  vector = Chroma(collection_name='customer', embedding_function=embeddings,
                  persist_directory='../chroma')
  vector.add_documents(documents=split_docs, embeddings=embeddings)


if __name__ == '__main__':
  web_search()
